import os
import asyncio
import tempfile
from contextlib import ExitStack
from typing import Text, Optional, List, Union, Dict

from rasa.importers.importer import TrainingDataImporter
from rasa import model
from rasa.model import FingerprintComparisonResult
from rasa.core.domain import Domain
from rasa.utils.common import TempDirectoryPath

from rasa.cli.utils import (
    print_success,
    print_warning,
    print_error,
    bcolors,
    print_color,
)
from rasa.constants import DEFAULT_MODELS_PATH, DEFAULT_CORE_SUBDIRECTORY_NAME, TMP_DIR


def train(
    domain: Text,
    config: Text,
    training_files: Union[Text, List[Text]],
    output: Text = DEFAULT_MODELS_PATH,
    force_training: bool = False,
    fixed_model_name: Optional[Text] = None,
    persist_nlu_training_data: bool = False,
    additional_arguments: Optional[Dict] = None,
    loop: Optional[asyncio.AbstractEventLoop] = None,
) -> Optional[Text]:
    """
    训练的总入口
    :param domain: 域yml文件：intents、responses、session_config
    :param config: 配置yml文件：language、pipeline、policies
    :param training_files:      训练文件：可以是文件夹
    :param output:              模型的输出路径
    :param force_training:      如果`True`则重新训练模型，即使数据未更改
    :param fixed_model_name:    要存储的模型名称
    :param persist_nlu_training_data:   `True` 则NLU训练数据应和model一起保存
    :param additional_arguments:        其他训练参数
    :param loop:
    :return:
    """
    if loop is None:
        try:
            # 创建事件循环（单线程：所有的函数调用都在loop中执行）
            loop = asyncio.get_event_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

    return loop.run_until_complete(
        train_async(
            domain=domain,
            config=config,
            training_files=training_files,
            output_path=output,
            force_training=force_training,
            fixed_model_name=fixed_model_name,
            persist_nlu_training_data=persist_nlu_training_data,
            additional_arguments=additional_arguments,
        )
    )


async def train_async(
    domain: Union[Domain, Text],
    config: Text,
    training_files: Optional[Union[Text, List[Text]]],
    output_path: Text = DEFAULT_MODELS_PATH,
    force_training: bool = False,
    fixed_model_name: Optional[Text] = None,
    persist_nlu_training_data: bool = False,
    additional_arguments: Optional[Dict] = None,
) -> Optional[Text]:
    """定义一个异步函数: 训练rasa模型 (Core and NLU).

    Args:
        domain: 域文件的路径
        config: Core和NLU的配置文件路径
        training_files: Core和NLU训练数据的路径
        output_path: 模型输出路径
        force_training: 如果`True`则重新训练模型，即使数据未更改
        fixed_model_name: 要存储的模型名称
        persist_nlu_training_data: `True` 则NLU训练数据应和model一起保存
        additional_arguments: 其他训练参数

    Returns:
        训练好的模型保存路径
    """

    # 给定输入参数调用 "训练数据输入器" 接口
    file_importer = TrainingDataImporter.load_from_config(
        config, domain, training_files)

    with ExitStack() as stack:

        train_path = stack.enter_context(TempDirectoryPath(tempfile.mkdtemp()))

        domain = await file_importer.get_domain()
        if domain.is_empty():
            return await handle_domain_if_not_exists(
                file_importer, output_path, fixed_model_name
            )

        return await _train_async_internal(
            file_importer,
            train_path,
            output_path,
            force_training,
            fixed_model_name,
            persist_nlu_training_data,
            additional_arguments,
        )


async def handle_domain_if_not_exists(
    file_importer: TrainingDataImporter, output_path, fixed_model_name
):
    nlu_model_only = await _train_nlu_with_validated_data(
        file_importer, output=output_path, fixed_model_name=fixed_model_name
    )
    print_warning(
        "Core training was skipped because no valid domain file was found. Only an nlu-model was created."
        "Please specify a valid domain using '--domain' argument or check if the provided domain file exists."
    )
    return nlu_model_only


async def _train_async_internal(
    file_importer: TrainingDataImporter,
    train_path: Text,
    output_path: Text,
    force_training: bool,
    fixed_model_name: Optional[Text],
    persist_nlu_training_data: bool,
    additional_arguments: Optional[Dict],
) -> Optional[Text]:
    """训练Rasa模型（Core和NLU）， 仅在`train_async`中使用.

    Args:
        file_importer:   提供训练数据（`TrainingDataImporter`）
        train_path:      训练模型的目录
        output_path:     输出路径
        force_training:  如果设为`True`，将重新训练模型，即使数据没有更改。
        persist_nlu_training_data:  如果该模型应保留NLU训练数据，则设置为“True”。
        fixed_model_name:           将要保存的模型的名称
        additional_arguments:       其他训练参数

    Returns:
        训练好的模型档案的路径
    """

    stories, nlu_data = await asyncio.gather(
        file_importer.get_stories(), file_importer.get_nlu_data()
    )

    if stories.is_empty() and nlu_data.is_empty():
        print_error(
            "No training data given. Please provide stories and NLU data in "
            "order to train a Rasa model using the '--data' argument."
        )
        return

    if stories.is_empty():
        print_warning("No stories present. Just a Rasa NLU model will be trained.")
        return await _train_nlu_with_validated_data(
            file_importer,
            output=output_path,
            fixed_model_name=fixed_model_name,
            persist_nlu_training_data=persist_nlu_training_data,
        )

    if nlu_data.is_empty():
        print_warning("No NLU data present. Just a Rasa Core model will be trained.")
        return await _train_core_with_validated_data(
            file_importer,
            output=output_path,
            fixed_model_name=fixed_model_name,
            additional_arguments=additional_arguments,
        )

    new_fingerprint = await model.model_fingerprint(file_importer)
    old_model = model.get_latest_model(output_path)
    fingerprint_comparison = FingerprintComparisonResult(force_training=force_training)
    if not force_training:
        fingerprint_comparison = model.should_retrain(
            new_fingerprint, old_model, train_path
        )

    if fingerprint_comparison.is_training_required():
        await _do_training(
            file_importer,
            output_path=output_path,
            train_path=train_path,
            fingerprint_comparison_result=fingerprint_comparison,
            fixed_model_name=fixed_model_name,
            persist_nlu_training_data=persist_nlu_training_data,
            additional_arguments=additional_arguments,
        )

        return model.package_model(
            fingerprint=new_fingerprint,
            output_directory=output_path,
            train_path=train_path,
            fixed_model_name=fixed_model_name,
        )

    print_success(
        "Nothing changed. You can use the old model stored at '{}'."
        "".format(os.path.abspath(old_model))
    )
    return old_model


async def _do_training(
    file_importer: TrainingDataImporter,
    output_path: Text,
    train_path: Text,
    fingerprint_comparison_result: Optional[FingerprintComparisonResult] = None,
    fixed_model_name: Optional[Text] = None,
    persist_nlu_training_data: bool = False,
    additional_arguments: Optional[Dict] = None,
):
    if not fingerprint_comparison_result:
        fingerprint_comparison_result = FingerprintComparisonResult()

    if fingerprint_comparison_result.should_retrain_core():
        await _train_core_with_validated_data(
            file_importer,
            output=output_path,
            train_path=train_path,
            fixed_model_name=fixed_model_name,
            additional_arguments=additional_arguments,
        )
    elif fingerprint_comparison_result.should_retrain_nlg():
        print_color(
            "Core stories/configuration did not change. "
            "Only the templates section has been changed. A new model with "
            "the updated templates will be created.",
            color=bcolors.OKBLUE,
        )
        await model.update_model_with_new_domain(file_importer, train_path)
    else:
        print_color(
            "Core stories/configuration did not change. No need to retrain Core model.",
            color=bcolors.OKBLUE,
        )

    if fingerprint_comparison_result.should_retrain_nlu():
        await _train_nlu_with_validated_data(
            file_importer,
            output=output_path,
            train_path=train_path,
            fixed_model_name=fixed_model_name,
            persist_nlu_training_data=persist_nlu_training_data,
        )
    else:
        print_color(
            "NLU data/configuration did not change. No need to retrain NLU model.",
            color=bcolors.OKBLUE,
        )


def train_core(
    domain: Union[Domain, Text],
    config: Text,
    stories: Text,
    output: Text,
    train_path: Optional[Text] = None,
    fixed_model_name: Optional[Text] = None,
    additional_arguments: Optional[Dict] = None,
) -> Optional[Text]:
    """训练core模型入口
    """
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(
        _train_core_async(
            domain=domain,
            config=config,
            stories=stories,
            output=output,
            train_path=train_path,
            fixed_model_name=fixed_model_name,
            additional_arguments=additional_arguments,
        )
    )


async def _train_core_async(
    domain: Union[Domain, Text],
    config: Text,
    stories: Text,
    output: Text,
    train_path: Optional[Text] = None,
    fixed_model_name: Optional[Text] = None,
    additional_arguments: Optional[Dict] = None,
) -> Optional[Text]:
    """训练 Core 模型

    参数:
        domain:     domain文件的路径
        config:     core文件的config配置路径
        stories:    Core training data的路径
        output:     输出路径
        train_path: 如果为“None”，则将在一个临时目录中训练模型，否则将在提供的目录中训练该模型。
        fixed_model_name:     要存储的模型名称
        uncompress:           如果为True，则不会压缩模型。
        additional_arguments: 其他训练参数

    返回:
        如果给出`train_path`，它将返回模型档案的路径，
        否则返回包含训练好的模型文件的目录的路径。

    """

    # 导入core训练的相关数据
    file_importer = TrainingDataImporter.load_core_importer_from_config(config, domain, [stories])

    # 获取域文件
    domain = await file_importer.get_domain()

    if domain.is_empty():
        print_error(
            "Core training was skipped because no valid domain file was found. "
            "Please specify a valid domain using '--domain' argument or check if the provided domain file exists."
        )
        return None

    # 获取story文件，进行验证
    if not await file_importer.get_stories():
        print_error(
            "No stories given. Please provide stories in order to "
            "train a Rasa Core model using the '--stories' argument."
        )
        return

    # 使用有效的数据进行core训练
    return await _train_core_with_validated_data(
        file_importer,
        output=output,
        train_path=train_path,
        fixed_model_name=fixed_model_name,
        additional_arguments=additional_arguments,
    )


async def _train_core_with_validated_data(
    file_importer: TrainingDataImporter,
    output: Text,
    train_path: Optional[Text] = None,
    fixed_model_name: Optional[Text] = None,
    additional_arguments: Optional[Dict] = None,
) -> Optional[Text]:
    """使用经过验证的训练数据和配置来训练core。"""

    import rasa.core.train

    with ExitStack() as stack:

        # 如果提供了train路径，请在出口处不进行任何操作。
        # 否则，请创建一条临时train路径，并在出口处进行清理。
        if train_path:
            _train_path = train_path
        else:
            _train_path = stack.enter_context(TempDirectoryPath(tempfile.mkdtemp()))

        # 正常（不比较）训练
        print_color("Training Core model...", color=bcolors.OKBLUE)
        domain, config = await asyncio.gather(
            file_importer.get_domain(), file_importer.get_config()
        )
        await rasa.core.train(
            domain_file=domain,
            training_resource=file_importer,
            output_path=os.path.join(_train_path, DEFAULT_CORE_SUBDIRECTORY_NAME),
            policy_config=config,
            additional_arguments=additional_arguments,
        )
        print_color("Core model training completed.", color=bcolors.OKBLUE)

        # 保存core模型
        if train_path is None:
            new_fingerprint = await model.model_fingerprint(file_importer)
            return model.package_model(
                fingerprint=new_fingerprint,
                output_directory=output,
                train_path=_train_path,
                fixed_model_name=fixed_model_name,
                model_prefix="core-",
            )

        return _train_path


def train_nlu(
    config: Text,
    nlu_data: Text,
    output: Text,
    train_path: Optional[Text] = None,
    fixed_model_name: Optional[Text] = None,
    persist_nlu_training_data: bool = False,
) -> Optional[Text]:
    """训练NLU模型。
    Args:
        config:     NLU的配置文件的路径。
        nlu_data:   NLU训练数据的路径。
        output:     模型详细信息输出路径。
        train_path: 如果设置为 “None”，则将在一个临时目录（temporary directory）中训练模型，否则在提供的目录中训练该模型。
        fixed_model_name: 要存储的模型的名称。
        persist_nlu_training_data: 如果要模型应保留NLU训练数据，则为Ture。
    Returns:
         如果给出`train_path`，它将返回模型文件的路径，否则返回包含训练好的模型文件的目录的路径。
    """

    # _train_nlu_async函数加入协程事件循环中
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(
        _train_nlu_async(
            config,
            nlu_data,
            output,
            train_path,
            fixed_model_name,
            persist_nlu_training_data,
        )
    )


async def _train_nlu_async(
    config: Text,
    nlu_data: Text,
    output: Text,
    train_path: Optional[Text] = None,
    fixed_model_name: Optional[Text] = None,
    persist_nlu_training_data: bool = False,
):
    # file_importer是导入的文件类：因为仅训练NLU模型，因此仍需选择训练文件。
    # 输出：importer的参数配置类，保存的是训练数据的所有importer，此时还没开始抽数据
    file_importer = TrainingDataImporter.load_nlu_importer_from_config(
        config, training_data_paths=[nlu_data]
    )

    # 训练数据类：此时的training_datas对象已经包含了所有的训练数据及各种实体、样例数量统计
    training_datas = await file_importer.get_nlu_data()

    if training_datas.is_empty():
        print_error(
            "No NLU data given. Please provide NLU data in order to train "
            "a Rasa NLU model using the '--nlu' argument."
        )
        return

    return await _train_nlu_with_validated_data(
        file_importer,
        output=output,
        train_path=train_path,
        fixed_model_name=fixed_model_name,
        persist_nlu_training_data=persist_nlu_training_data,
    )


async def _train_nlu_with_validated_data(
    file_importer: TrainingDataImporter,
    output: Text,
    train_path: Optional[Text] = None,
    fixed_model_name: Optional[Text] = None,
    persist_nlu_training_data: bool = False,
) -> Optional[Text]:
    """使用经过验证的训练和配置数据训练NLU。
    真正的训练是从"rasa/nlu/train"开始
    """

    import rasa.nlu.train

    with ExitStack() as stack:
        if train_path:
            # 如果提供了训练路径，退出时不做任何处理
            _train_path = train_path
        else:
            # 否则，新建临时的训练路径，并且在退出时进行清理
            _train_path = stack.enter_context(TempDirectoryPath(tempfile.mkdtemp()))

        # 得到配置路径，
        config = await file_importer.get_config()

        # NLU模型开始训练，需要输入config字典、file_importer对象、临时训练路径
        print_color("Training NLU model...", color=bcolors.OKBLUE)
        _, nlu_model, _ = await rasa.nlu.train(
            config,
            file_importer,
            _train_path,
            fixed_model_name="nlu",
            persist_nlu_training_data=persist_nlu_training_data,
        )
        print_color("NLU model training completed.", color=bcolors.OKBLUE)

        # 仅仅 NLU 被训练
        if train_path is None:

            # 创建模型指纹
            new_fingerprint = await model.model_fingerprint(file_importer)

            # 返回模型路径
            return model.package_model(
                fingerprint=new_fingerprint,
                output_directory=output,
                train_path=_train_path,
                fixed_model_name=fixed_model_name,
                model_prefix="nlu-",
            )

        return _train_path
